// Copyright 2020-2024 CesiumGS, Inc. and Contributors

#pragma once

#include "CoreMinimal.h"

#if WITH_EDITOR
#include "Materials/MaterialFunctionMaterialLayer.h"
#endif

#include <memory>

#if WITH_EDITOR
enum ECustomMaterialOutputType : int;

enum class ECesiumEncodedMetadataType : uint8;
enum class ECesiumEncodedMetadataComponentType : uint8;

class UMaterialExpressionParameter;
#endif

namespace GenerateMaterialUtility {
#if WITH_EDITOR
template <typename ObjClass>
static FORCEINLINE ObjClass* LoadObjFromPath(const FName& Path) {
  if (Path == NAME_None)
    return nullptr;

  return Cast<ObjClass>(
      StaticLoadObject(ObjClass::StaticClass(), nullptr, *Path.ToString()));
}

/**
 * @brief The message used for the description of autogenerated nodes.
 */
static const FString AutogeneratedMessage = "AUTOGENERATED DO NOT EDIT";

/**
 * @brief An increment constant that is used to space out the autogenerated
 * nodes.
 */
static const int32 Incr = 200;

/**
 * @brief Classification of nodes into categories, such that they can be handled
 * separately.
 */
struct MaterialNodeClassification {
  /** Nodes that were procedurally created with the AutogeneratedMessage. */
  TArray<UMaterialExpression*> AutoGeneratedNodes;
  /** Nodes without the AutogeneratedMessage (presumably added by the user). */
  TArray<UMaterialExpression*> UserAddedNodes;
};

/**
 * @brief A container that bundles various components used to procedurally
 * create or modify a material.
 */
struct MaterialGenerationState {
  /**
   * Nodes that are procedurally created and to be marked with
   * AutogeneratedMessage.
   */
  TArray<UMaterialExpression*> AutoGeneratedNodes;
  /**
   * Nodes that should only exist once in the material layer, e.g.,
   * SetMaterialAttributes. These will not be regenerated if already present.
   */
  TArray<UMaterialExpression*> OneTimeGeneratedNodes;
  /**
   * A map from autogenerated nodes to the FExpressionInputs that they
   * previously sent outputs to.
   */
  TMap<FString, TArray<FExpressionInput*>> ConnectionOutputRemap;
  /**
   * A map from autogenerated nodes to the FExpressionInputs that they
   * previously took inputs from.
   */
  TMap<FString, TMap<FString, const FExpressionInput*>> ConnectionInputRemap;
};

/**
 * Creates a new material layer asset with the given name. If the specified
 * package doesn't exist, this function creates it.
 */
UMaterialFunctionMaterialLayer*
CreateMaterialLayer(const FString& PackageName, const FString& MaterialName);

/**
 * Moves the generated nodes from the material state into the given material
 * layer.
 */
void MoveNodesToMaterialLayer(
    MaterialGenerationState& MaterialState,
    UMaterialFunctionMaterialLayer* pMaterialLayer);

/**
 * Computes a scalar for spacing out material nodes. The actual computation is
 * rather arbitrary, but this prevents clumping when properties have extremely
 * long names.
 */
float GetNameLengthScalar(const FName& Name);
float GetNameLengthScalar(const FString& Name);

ECustomMaterialOutputType
GetOutputTypeForEncodedType(ECesiumEncodedMetadataType Type);

FString GetHlslTypeForEncodedType(
    ECesiumEncodedMetadataType Type,
    ECesiumEncodedMetadataComponentType ComponentType);

FString GetSwizzleForEncodedType(ECesiumEncodedMetadataType Type);

/**
 * @brief Generates a parameter node corresponding to the given encoded
 * metadata type.
 */
UMaterialExpressionParameter* GenerateParameterNode(
    UMaterialFunctionMaterialLayer* TargetMaterialLayer,
    const ECesiumEncodedMetadataType Type,
    const FString& Name,
    int32 NodeX,
    int32 NodeY);
#endif

} // namespace GenerateMaterialUtility
